# -*- coding: utf-8 -*-
# @Author: Tim Liu
# @Date: 2024-05-25
# @Last Modified by: Tim Liu
# @Last Modified time: 2024-05-25

# @Description: Knowledge Base DAL for RAG

from sqlalchemy.ext.asyncio import AsyncSession
from core.crud import DalBase
from crewplus.apps.rag import models, schemas


class DocumentDal(DalBase):

    def __init__(self, db: AsyncSession):
        super(DocumentDal, self).__init__()
        self.db = db
        self.model = models.DocumentDB
        self.schema = schemas.DocumentSimpleOut

    async def add_view_number(self, data_id: int) -> None:
        """
        update document's view number +1
        """
        obj: models.DocumentDB = await self.get_data(data_id)
        obj.view_number = obj.view_number + 1 if obj.view_number else 1
        await self.flush(obj)


class KnowledgeBaseDal(DalBase):

    def __init__(self, db: AsyncSession):
        super(KnowledgeBaseDal, self).__init__()
        self.db = db
        self.model = models.KnowledgeBaseDB
        self.schema = schemas.KnowledgeBaseSimpleOut

    # def get_id_by_coll_id_and_name(self, coll_id: int, coll_name: str):
    #     query = select(models.KnowledgeBaseDB.id).where(
    #         models.KnowledgeBaseDB.coll_id == coll_id,
    #         models.KnowledgeBaseDB.coll_name == coll_name
    #     )
    #     result = await self.db.execute(query)
    #     data = result.scalar_one_or_none()
    #     return data
